import rrd_base
import re
import time

import snmp
import color

try:
    from pyrrd.rrd import DataSource, RRA, RRD
    from pyrrd.graph import DEF, CDEF, VDEF
    from pyrrd.graph import LINE, AREA, GPRINT
    from pyrrd.graph import ColorAttributes, Graph
except:
    print "Couldn't load required modules. You probably need to install it with:"
    sys.exit('sudo apt-get install python-pyrrd python-pysnmp4')

class RrdNet(rrd_base.RrdBase):
    def __init__(self, domain, name, alt):
        super(RrdNet,self).__init__(domain=domain, group="net", instance=name)
        self.alt = alt

    def appendDataSources(self, dataSources):
        assert 0, "RrdNet: appendDataSources() should not be called"

    def create_file(self):
        dataSources = []
        self.appendDataSources(dataSources)

        roundRobinArchives = []
        for func in ['AVERAGE', 'MAX', 'MIN']:
            for step,row in self.rra_pattern():
                roundRobinArchives.append(RRA(cf=func, xff=0.5, steps=step, rows=row))

        self.myRRD = RRD(self.filename, ds=dataSources, rra=roundRobinArchives, step=self.data_step())
        self.myRRD.create()

    def graph(self):
        def_rxB       = DEF(rrdfile=self.myRRD.filename, vname='d0', cdef="AVERAGE", dsName='rx_bytes')
        def_txB       = DEF(rrdfile=self.myRRD.filename, vname='d1', cdef="AVERAGE", dsName='tx_bytes')
        def_rxB_max   = DEF(rrdfile=self.myRRD.filename, vname='d2', cdef="MAX",     dsName='rx_bytes')
        def_txB_max   = DEF(rrdfile=self.myRRD.filename, vname='d3', cdef="MAX",     dsName='tx_bytes')
        cdef_rxb      = CDEF(vname="c0",          rpn="{0},8,*".format(def_rxB.vname))
        cdef_rxb_max  = CDEF(vname="c1",          rpn="{0},8,*".format(def_rxB_max.vname))
        cdef_txb      = CDEF(vname="c2",          rpn="{0},8,*".format(def_txB.vname))
        cdef_txb_max  = CDEF(vname="c3",          rpn="{0},8,*".format(def_txB_max.vname))
        vdef_rxb_p    = VDEF(vname="v0",          rpn="{0},MAXIMUM".format(cdef_rxb_max.vname))
        vdef_rxB_total= VDEF(vname="rTotal",      rpn="{0},TOTAL".format(def_rxB.vname))
        vdef_txB_total= VDEF(vname="tTotal",      rpn="{0},TOTAL".format(def_txB.vname))
        vdef_txb_p    = VDEF(vname="v1",          rpn="{0},MAXIMUM".format(cdef_txb_max.vname))
        area_rx       = AREA(defObj=cdef_rxb,       color=color.RGBA['red'], legend="RX")
        area_tx       = AREA(defObj=cdef_txb,       color=color.RGBA['blue'],  legend="TX")
        line_0        = LINE(1, value="1024",       color=color.RGB['black'])
        line_1        = LINE(1, value="10240",      color=color.RGB['black'])
        line_2        = LINE(1, value="102400",     color=color.RGB['black'])
        line_3        = LINE(1, value="1048576",    color=color.RGB['black'])
        line_4        = LINE(1, value="10485760",   color=color.RGB['black'])
        l_tx_max      = LINE(1, defObj=cdef_txb_max, color=color.RGBA['red'])
        l_rx_max      = LINE(1, defObj=cdef_rxb_max, color=color.RGBA['red'])
        l_tx_peak     = LINE(2, defObj=vdef_txb_p,  color=color.RGBA['black'])
        l_rx_peak     = LINE(2, defObj=vdef_rxb_p,  color=color.RGBA['black'])
        p_tx_peak     = GPRINT(vdef_txb_p,  "Peak Tx %.2lf%sb/s")
        p_rx_peak     = GPRINT(vdef_rxb_p,  "Peak Rx %.2lf%sb/s")
        p_txB_total   = GPRINT(vdef_txB_total, "Total Tx %.1lf %sB")
        p_rxB_total   = GPRINT(vdef_rxB_total, "Total Rx %.1lf %sB")

        for tf in self.graph_time_frames().keys():
            g = Graph(self.image_name(tf),
                      start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                      base=1024, logarithmic=True,
                      title='"Network Traffic for {0}/{1}{3} - Updated {2}"'.format(self.data_domain, self.data_instance, time.strftime("%Y%m%d at %H:%M"), self.alt),
                      vertical_label='Bits/Second',
                      height=self.graph_height, width=self.graph_width,
                      x_grid=self.grid[tf],
                      step=self.graph_step_times[tf])
            g.data.extend([def_rxB, def_txB, def_rxB_max, def_txB_max, cdef_rxb, cdef_txb, cdef_rxb_max, cdef_txb_max, vdef_rxb_p, vdef_txb_p, area_rx, p_rx_peak, area_tx, p_tx_peak, vdef_rxB_total, vdef_txB_total, p_txB_total, p_rxB_total, line_0, line_1, line_2, line_3, line_4 ])

            if self.graph_time_frames()[tf] > self.graph_time_frames()["1day"] and self.graph_time_frames()[tf] < self.graph_time_frames()["1month"]:
                g.data.extend([l_tx_max, l_rx_max, l_tx_peak, l_rx_peak])

            g.write()

    def update(self):
        pass

class RrdNetLocal(RrdNet):
    def __init__(self, domain, name, alt):
        super(RrdNetLocal,self).__init__(domain=domain, name=name, alt=alt)

    def appendDataSources(self, dataSources):
        for ds in ["rx_bytes", "rx_packets", "rx_errs", "rx_drop", "tx_bytes", "tx_packets", "tx_errs", "tx_drop" ]:
            dataSources.append(DataSource( dsName=ds, dsType='COUNTER', heartbeat=self.heartbeat()))

    def update(self):
        f = open("/proc/net/dev")
        for line in f:
            if re.search(self.data_instance, line):
                words = line.split()
                self.myRRD.bufferValue(int(round(time.time())), words[1], words[2], words[3], words[4], words[9], words[10], words[11], words[12])

                self.myRRD.update()
                break

class RrdNetSnmp(RrdNet):
    def __init__(self, domain, name, alt):
        super(RrdNetSnmp,self).__init__(domain=domain, name=name, alt=alt)

    def getDdwrtSnmpIndex(self, n):
        l = {  "vlan0": 5
              ,"vlan1": 6
              ,"br0":   7
              ,"eth1":  4
            }
        return l[n]

    def appendDataSources(self, dataSources):
        for ds in ["rx_bytes", "tx_bytes"]:
            dataSources.append(DataSource( dsName=ds, dsType='COUNTER', heartbeat=self.heartbeat()))

    def update(self):
#                   _________
#                  (___br0___) 
#                   /  | \ \_ \____ 
#                  /   |  \  \__   \_____ 
#                 /    |   \    \        \ 
#      +----------|----|----|----|--+ +---|----+ 
#      | [0]     [1]  [2]  [3]  [4] | | [wlan] | 
#      |vlan1   ----- vlan0 ------  | |        | 
#      |  WAN           LAN         | |  WLAN  | 
#      |                            | |        | 
#      |Internet   1    2    3    4 | |  WLAN  | <-- label 
#      +----------------------------+ +--------+ 
#            ethernet switch           wireless 
#           eth0 (cpu port 5)            eth1
#
#iso.3.6.1.2.1.2.2.1.2.1  = STRING: "lo"
#iso.3.6.1.2.1.2.2.1.2.2  = STRING: "teql0"
#iso.3.6.1.2.1.2.2.1.2.3  = STRING: "eth0"
#iso.3.6.1.2.1.2.2.1.2.4  = STRING: "eth1"
#iso.3.6.1.2.1.2.2.1.2.5  = STRING: "vlan0"
#iso.3.6.1.2.1.2.2.1.2.6  = STRING: "vlan1"
#iso.3.6.1.2.1.2.2.1.2.7  = STRING: "br0"
#iso.3.6.1.2.1.2.2.1.2.8  = STRING: "etherip0"
#iso.3.6.1.2.1.2.2.1.2.9  = STRING: "imq0"
#iso.3.6.1.2.1.2.2.1.2.10 = STRING: "imq1"
        snmpIndex = self.getDdwrtSnmpIndex(self.data_instance)
        tx_bytes = snmp.get((1,3,6,1,2,1,2,2,1,16, snmpIndex))
        rx_bytes = snmp.get((1,3,6,1,2,1,2,2,1,10, snmpIndex))

        self.myRRD.bufferValue(int(round(time.time())), int(rx_bytes), int(tx_bytes))
        self.myRRD.update()
